//
// Created by Alex on 17/3/14.
//

#ifndef LIBCR_CONNECTION_H
#define LIBCR_CONNECTION_H

#include <stack>
#include <thread>
#include <random>
#include "socket.h"
#include "buffer.h"

struct Connection;
struct ConnectionImpl;
typedef void (*HandlePacket)(const Connection &conn, const void *data, size_t size);

struct Connection {
  void SendData(const void *data, size_t size) {
    m_impl->SendData(data, size);
  }

  std::string GetIP() {
    struct in_addr addr;
    addr.s_addr = m_ip;
    return inet_ntoa(addr);
  }

  int GetPort() {
    return m_port;
  }

  uint32_t m_ip;
  int m_port;
  ConnectionImpl *m_impl;
};

struct PacketTask {
  Connection m_conn;
  void *m_data;
  size_t m_size;
  HandlePacket m_callback;
  PacketTask() : m_data(nullptr), m_size(0), m_callback(nullptr) {}
};

struct Worker {

  Worker(WorkerPool *pool) :
      m_thread(nullptr),
      m_taskCount(0),
      m_shutdown(false),
      m_exit(false),
      m_workers(pool) {}

  virtual ~Worker() {
    if (m_thread == nullptr) {
      return;
    }

    while (1) {
      sleep(1);

      if (m_exit) {
        delete m_thread;
        m_thread = nullptr;
        break;
      }

    }
  }

  void Run() {
    for (int i = 0; i < m_maxUthreadCnt; ++i) {
      cr_go [m_taskChannel] {

        while (true) {
          PacketTask task;
          m_taskChannel >> task;

          if (task.m_data == nullptr) {
            cr_sleep(100);
            break;
          }

          HandlePacket callback = task.m_callback;
          if (callback == nullptr) {
            continue;
          }

          callback(task.m_conn, task.m_data, task.m_size);
          delete task.m_data;
        }
      };
    }

    Worker *worker = this;
    cr_go [worker] {
      while (!worker->m_shutdown || !worker->m_tasks.empty()) {
        std::unique_lock lock(worker->m_lock);
        PacketTask task;

        //没有数据 sleep 200ms
        if (worker->m_tasks.empty()) {
          worker->m_cond.wait_until(lock, std::chrono::milliseconds(200));
        }

        //200ms后, 有数据: 取自己的task, 没有: 取别的worker的task
        if (worker->m_tasks.empty()) {
          lock.unlock();
          PickTaskFromWorkers(task);

        } else {
          task = worker->m_tasks.front();
          worker->m_tasks.pop_front();
          lock.unlock();
        }

        //回调
        if (task.m_data != nullptr) {
          worker->m_taskChannel << task;
        }
      }
    };

    cr_scheduler();
    m_exit = true;
  }

  void AddTask(const PacketTask &task) {
    if (m_shutdown) {
      return;
    }

    m_lock.lock();

    m_tasks.push_back(task);
    m_taskCount += 1;

    m_lock.unlock();

    m_cond.notify_one();
  }

  void Shutdown() {
    m_shutdown = true;
  }

  void Start() {
    m_thread = new std::thread(Worker::Run());
    m_thread->detach();
  }

  bool PickTaskFromWorkers(PacketTask &task) {
    if (m_thread == nullptr || m_workers == nullptr) {
      return false;
    }

    if (m_workers->m_workers.size() == 1) {
      return false;
    }

    std::uniform_int_distribution<int> dist(0, m_workers->m_workers.size());

    int retry = 0;
    while (true) {
      size_t idx = dist(m_rand) % m_workers->m_workers.size();
      Worker *worker = m_workers->m_workers[idx];

      if (worker->m_shutdown || worker->m_thread->get_id() == m_thread->get_id()) {
        continue;
      }

      worker->m_lock.lock();

      if (worker->m_tasks.empty()) {
        worker->m_lock.unlock();

        if (++retry == m_workers->m_workers.size() / 2) {
          return false;
        } else {
          continue;
        }
      }

      task = worker->m_tasks.front();
      worker->m_tasks.pop_front();
      worker->m_taskCount -= 1;

      worker->m_lock.unlock();

      return true;
    }
  };

private:
  bool m_shutdown;
  bool m_exit;
  std::thread *m_thread;
  std::list<PacketTask> m_tasks;
  int m_taskCount;
  std::mutex m_lock;
  std::condition_variable m_cond;

  WorkerPool *m_workers;
  std::random_device m_rand;

  size_t m_maxUthreadCnt;
  Channel<PacketTask> m_taskChannel;
};

struct WorkerPool {

  WorkerPool(size_t cnt) : m_nextWorker(0) {
    for (int i = 0; i < cnt; ++i) {
      m_workers.push_back(new Worker(this));
    }
  }

  void AddTask(const PacketTask &task) {
    ssize_t idx = __sync_add_and_fetch(&m_nextWorker, 1) % m_workers.size();
    Worker *worker = m_workers[idx];
    worker->AddTask(task);
  }

  void Shutdown() {
    for (Worker *worker : m_workers) {
      worker->Shutdown();
    }
  }

  void Start() {
    for (Worker *worker : m_workers) {
      worker->Run();
    }

    for (Worker *worker : m_workers) {
      delete worker;
    }

    m_workers.clear();
  }

  std::vector<Worker *> m_workers;
  int m_nextWorker;
};

struct InputBufferProcessor : public PacketProcessor {
  InputBufferProcessor(uint16_t magic, ConnectionImpl *impl, WorkerPool *workers) :
      PacketProcessor(magic),
      m_workers(workers) {

    m_conn.m_ip = inet_addr(impl->m_socket->m_peerIp.c_str());
    m_conn.m_port = impl->m_socket->m_peerPort;
    m_conn.m_impl = impl;
  }

  void SetPacketCallback(HandlePacket callback) {
    m_callback = callback;
  }

  virtual void HandlerPacket(PacketHeader *header, const char *data, size_t len) {
    if (data != nullptr) {

      if (m_workers != nullptr) {
        char *copy = new char[len];
        memcpy(copy, data, len);

        PacketTask task;
        task.m_size = len;
        task.m_data = (void *) data;
        task.m_conn = m_conn;
        task.m_callback = m_callback;

        m_conn.m_impl->Ref();
        m_workers->AddTask(task);

      } else if (m_callback != nullptr) {

        m_callback(m_conn, (void *) data, len);
      }

    }
  }

private:
  WorkerPool *m_workers;
  Connection m_conn;
  HandlePacket m_callback;
};

struct ConnectionImpl {
  ConnectionImpl(WorkerPool *pool) {
    m_inputProcessor = new InputBufferProcessor(m_magic, this, pool);
    m_outputProcessor = new OutputBufferProcessor(m_magic);
    m_outputBufferSize = 0;
  }

  void SetCallback(HandlePacket callback) {
    if (m_inputProcessor != nullptr) {
      m_inputProcessor->SetPacketCallback(callback);
    }
  }

  virtual ~ConnectionImpl() {
    if (m_socket != nullptr) {
      m_socket->UnRef();
    }
    delete m_inputProcessor;
    delete m_outputProcessor;

    for (std::pair<void *, size_t> &task : m_outputBuffer) {
      delete task.first;
    }
    m_outputBuffer.clear();
  }

  void Ref() {
    m_refCnt += 1;
  }

  bool UnRef() {
    m_refCnt -= 1;
    if (m_refCnt <= 0) {
      delete this;
      return true;
    }
    return false;
  }

  void SendData(const void *data, size_t size) {
    if (data == nullptr || size <= 0) {
      return;
    }
    char *buffer = new char[size];
    memcpy(buffer, data, size);

    pthread_spin_lock(&m_lock);

    m_outputBuffer.push_back(std::make_pair(buffer, size));
    m_outputBufferSize += 1;
    cr_cond_signal(&m_outputCond);

    pthread_spin_unlock(&m_lock);
  }

  void OnRead() {
    //接收协程
    ConnectionImpl *conn = this;
    cr_go [conn] {
      while (true) {
        ssize_t ret = 0;
        int size = conn->m_inputBuffer.RecvToBuffer(
            conn->m_socket->m_socketFd, ret);

        if (ret <= 0) {
          break;
        }

        if (size > 0) {
          conn->m_inputProcessor->OnReadComplete(conn->m_inputBuffer);
        }
      }

      conn->UnRef();
    };
  }

  void OnWrite() {
    //发送协程
    ConnectionImpl *conn = this;
    cr_go [conn] {
      while (true) {
        std::pair<void *, size_t> buffer = std::make_pair(nullptr, 0);

        pthread_spin_lock(&conn->m_lock);
        if (conn->m_outputBufferSize <= 0) {

          pthread_spin_unlock(&conn->m_lock);
          cr_cond_timedwait(&(conn->m_outputCond), 0);
          pthread_spin_lock(&conn->m_lock);

        }

        if (conn->m_outputBufferSize > 0) {
          buffer = conn->m_outputBuffer.front();
          conn->m_outputBuffer.pop_front();
        }

        pthread_spin_unlock(&conn->m_lock);

        if (buffer.first != nullptr && buffer.second > 0) {

          ssize_t ret = conn->m_outputProcessor->WriteData(
              conn->m_socket->m_socketFd, 0, buffer.first, buffer.second);

          if (ret <= 0) {
            break;
          }

        }
      }
      conn->UnRef();
    };

  }

  uint16_t m_magic;
  Socket *m_socket;

  BufferList m_inputBuffer;
  InputBufferProcessor *m_inputProcessor;
  OutputBufferProcessor *m_outputProcessor;

  pthread_spinlock_t m_lock;
  std::list<std::pair<void *, size_t> > m_outputBuffer;
  size_t m_outputBufferSize;
  cr_cond_t m_outputCond;

  std::atomic_int m_refCnt;
};

struct IOThread {
  IOThread(size_t workerCnt) {
    m_workers = new WorkerPool(workerCnt);
  }

  void AddConnection(Socket *socket) {
    uint64_t data = (uint64_t) socket;
    cr_write(m_eventFd, &data, sizeof(data));
  }

  void SetPacketCallback(HandlePacket callback) {
    m_callback = callback;
  }

  void Run() {
    m_eventFd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC);
    if (m_eventFd < 0 && errno == EINVAL) {
      m_eventFd = eventfd(0, 0);
    }

    if (m_eventFd == -1) {
      exit(0);
    }

    m_workers->Start();
    m_ioThread = new std::thread(LoopInner);
    m_ioThread->detach();
  }

  void LoopInner() {
    cr_go [m_eventFd] {
      uint64_t data;
      while (true) {
        ssize_t ret = cr_read(m_eventFd, &data, sizeof(data));
        if (ret < 0) {
          continue;
        }
        if (data <= 0) {
          break;
        }
        Socket *socket = (Socket *) data;
        AddConnectionInner(socket);
      }
    };

    cr_scheduler();
  }

  void AddConnectionInner(Socket *socket) {
    ConnectionImpl *conn = new ConnectionImpl(m_workers);
    conn->SetCallback(m_callback);
    conn->m_socket = socket;

    conn->Ref();
    conn->Ref();

    conn->OnRead();
    conn->OnWrite();
  }

  HandlePacket m_callback;
  WorkerPool *m_workers;
  std::thread *m_ioThread;
  int m_eventFd;
};

struct TCPServer {
  TCPServer(const std::string &ip, int port) {
    m_listenSocket = new ServerSocket(ip, port);
    m_ioThreadCnt = 4;
    m_workerCnt = m_ioThreadCnt * 2;
  }

  void SetPacketCallback(HandlePacket callback) {
    m_callback = callback;
  }

  void SetThread(size_t ioThreadCnt, size_t workerCnt) {
    m_ioThreadCnt = ioThreadCnt;
    m_workerCnt = workerCnt;
  }

  void Start() {
    int perWorkerCnt = 2;
    while (perWorkerCnt * m_ioThreadCnt < m_workerCnt) {
      perWorkerCnt += 1;
    }

    for (int i = 0; i < m_ioThreadCnt; ++i) {
      IOThread *thread = new IOThread(perWorkerCnt);
      thread->SetPacketCallback(m_callback);
      thread->Run();
      m_ioThreads.push_back(thread);
    }

    TCPServer *server = this;
    m_listenSocket->Listen();

    cr_go [server] {
      while (true) {
        Socket *cli = server->m_listenSocket->Accept();
        if (cli == nullptr) {
          return;
        }
        std::uniform_int_distribution<int> dist(0, m_ioThreadCnt);;
        IOThread *thread = m_ioThreads[dist(m_rand) % m_ioThreadCnt];
        thread->AddConnection(cli);
      }
    };

  }

  HandlePacket m_callback;
  ServerSocket *m_listenSocket;
  size_t m_ioThreadCnt;
  size_t m_workerCnt;
  std::vector<IOThread *> m_ioThreads;
  std::random_device m_rand;
};

struct TCPClient {

  TCPClient() : m_workerCnt(0) {}

  void SetWorkers(size_t worker) {
    m_workerCnt = worker;
  }

  void SetPacketCallback(HandlePacket callback) {
    m_callback = callback;
  }

  void SetReConnectInterval(long ms);

  bool Connect(const std::string &connStr) {
    m_socket = new Socket(connStr);
    bool open = m_socket->Open();
    if (open) {
      ConnectionImpl *conn = new ConnectionImpl(nullptr);
      m_conn = conn;

      conn->SetCallback(m_callback);
      conn->m_socket = m_socket;

      conn->Ref();
      conn->Ref();
      conn->OnRead();
      conn->OnWrite();
    }

    return open;
  }

  void Send(const char *data, size_t size) {
    if (m_conn != nullptr) {
      m_conn->SendData(data, size);
    }
  }

  void Close() {
    m_socket->Close();
  }

private:
  HandlePacket m_callback;
  ConnectionImpl *m_conn;
  Socket *m_socket;
  size_t m_workerCnt;
};

void Client() {
  TCPClient *client = new TCPClient;
  client->SetPacketCallback(nullptr);

  for (int i = 0; i < 100; ++i) {
    client->Send("dddd", 3);
  }

  cr_scheduler();
}

void Server() {
  TCPServer *server = new TCPServer(nullptr, 9098);
  server->SetPacketCallback(nullptr);
  server->Start();
  cr_scheduler();
}

int main() {
  if (fork()) {
    Client();
  } else {
    Server();
  }
}

#endif //LIBCR_CONNECTION_H
